#include "tcp_sender.hh"

#include "tcp_config.hh"

#include <algorithm>
#include <iostream>
#include <random>

using namespace std;

//! \param[in] capacity the capacity of the outgoing byte stream
//! \param[in] retx_timeout the initial amount of time to wait before retransmitting the oldest outstanding segment
//! \param[in] fixed_isn the Initial Sequence Number to use, if set (otherwise uses a random ISN)
TCPSender::TCPSender(const size_t capacity, const uint16_t retx_timeout, const std::optional<WrappingInt32> fixed_isn)
    : _stream(capacity)
    , _initial_retransmission_timeout(retx_timeout)
    , _isn(fixed_isn.value_or(WrappingInt32{random_device()()})) {}

uint64_t TCPSender::bytes_in_flight() const {
    uint64_t count{0};
    auto i = _outstanding_segments.begin();
    auto e = _outstanding_segments.end();
    while (i != e) {
        count += i->length_in_sequence_space();
        ++i;
    }
    return count;
}

void TCPSender::fill_window() {
    while (_window_size > 0) {
        if (_fin_sent) {
            // should send no more segments
            return;
        }
        TCPSegment segment{};
        if (_next_seqno == 0) {
            // set SYN flag for the first segment
            segment.header().syn = 1;
        } else if (_stream.buffer_empty() && !_stream.input_ended()) {
            break;
        }
        // fill header and payload
        segment.header().seqno = wrap(_next_seqno, _isn);
        size_t read_size =
            min(min(static_cast<size_t>(_window_size), _stream.buffer_size()), TCPConfig::MAX_PAYLOAD_SIZE);
        segment.payload() = Buffer{std::move(_stream.read(read_size))};
        // set FIN flag for the last segment
        if (_stream.input_ended() && _stream.buffer_empty() && (_window_size > read_size)) {
            segment.header().fin = true;
            _fin_sent = true;
        }
        _next_seqno += segment.length_in_sequence_space();
        _window_size -= segment.length_in_sequence_space();
        // send segment and buffer it
        _segments_out.push(segment);
        _outstanding_segments.push_back(segment);
        // start timer if the timer is off
        if (!_timer.started()) {
            _timer.start();
        }
    }
}

//! \param ackno The remote receiver's ackno (acknowledgment number)
//! \param window_size The remote receiver's advertised window size
void TCPSender::ack_received(const WrappingInt32 ackno, const uint16_t window_size) {
    // error: the acked seqno hasn't been sent yet
    if (ackno.raw_value() > wrap(_next_seqno, _isn).raw_value()) {
        return;
    }
    // Is there new acked segments?
    bool new_acked{false};
    // determine the behaviour when the timer is expired
    _last_window_size = window_size;
    // see window size 0 as 1
    uint16_t new_window_size = (window_size == 0 ? 1 : window_size);
    // remove fully acked segments
    while (!_outstanding_segments.empty()) {
        auto front = _outstanding_segments.front();
        if (front.header().seqno.raw_value() + static_cast<uint32_t>(front.length_in_sequence_space()) <=
            ackno.raw_value()) {
            // remove fully acknowledged segments
            _outstanding_segments.pop_front();
            _consecutive_retransmissions_num = 0;
            new_acked = true;
        } else {
            // if there's still segment not fully acknowledged
            break;
        }
    }
    // timer's behaviour
    if (new_acked) {
        _timer.stop();
        _timer.set_RTO(_initial_retransmission_timeout);
        if (!_outstanding_segments.empty()) {
            _timer.start();
        }
    }
    // update window size
    if (unwrap(ackno, _isn, _next_seqno) + new_window_size > _next_seqno + _window_size) {
        _window_size = unwrap(ackno, _isn, _next_seqno) + new_window_size - _next_seqno;
    }
}

//! \param[in] ms_since_last_tick the number of milliseconds since the last call to this method
void TCPSender::tick(const size_t ms_since_last_tick) {
    if (!_timer.started()) {
        return;
    }
    // timer count down
    _timer.tick(ms_since_last_tick);
    if (_timer.expired()) {
        _segments_out.push(_outstanding_segments.front());
        if (_last_window_size != 0) {
            ++_consecutive_retransmissions_num;
            _timer.set_RTO(2 * _timer.get_RTO());
        }
        _timer.start();
    }
}

void TCPSender::send_empty_segment() {
    TCPSegment segment{};
    segment.header().seqno = wrap(_next_seqno, _isn);
    _segments_out.push(segment);
}

/*
 * TCPTimer function implementations
 */
TCPTimer::TCPTimer(const unsigned int timeout) : _retransmission_timeout(timeout) {}

void TCPTimer::start() {
    _started = true;
    _countdown = _retransmission_timeout;
}

void TCPTimer::stop() {
    _started = false;
    _countdown = 0;
}

void TCPTimer::tick(const size_t ms) { _countdown -= min(ms, static_cast<size_t>(_countdown)); }